import random
import os  
import torch  
from torch.utils.data import Dataset, DataLoader  
from torchvision import transforms  
from PIL import Image  
from collections import defaultdict  
import json
  
class CustomImageDataset(Dataset):  
    def __init__(self, root_dir, transform=None, train=True, split_ratio=0.8):  
        self.root_dir = root_dir  
        self.transform = transform  
        self.train = train  
        self.split_ratio = split_ratio  
    
        self.class_to_samples = defaultdict(list)  
        for cls_name in os.listdir(root_dir):  
            cls_dir = os.path.join(root_dir, cls_name)  
            if os.path.isdir(cls_dir):  
                for file in os.listdir(cls_dir):  
                    if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):  
                        path = os.path.join(cls_dir, file)  
                        self.class_to_samples[cls_name].append(path)  
        
        random.seed(123) # 保证每次实验划分一致
        self.samples = []  
        for cls_name, samples in self.class_to_samples.items():  
            random.shuffle(samples)  
            split_index = int(len(samples) * split_ratio)  
            if train:  
                self.samples.extend([(path, cls_name) for path in samples[:split_index]])  
            else:  
                self.samples.extend([(path, cls_name) for path in samples[split_index:]])  
                
  
        self.classes = list(self.class_to_samples.keys())  
        self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)}  
        
        if self.train is False:
            self.save_class_labels('class_labels.json')
        
    def save_class_labels(self, json_file):
        with open(json_file, 'w') as f:
            json.dump(self.class_to_idx, f, indent=4)
    
    
    def __len__(self):  
        return len(self.samples)  
  
    def __getitem__(self, idx):   
        path, cls_name = self.samples[idx]  
        image = Image.open(path).convert('RGB')  
        label = self.class_to_idx[cls_name]   
        if self.transform:  
            image = self.transform(image)  
        return image, label  